from src import Preproduction as Pre

id, ch = Pre.load("results/hanzi_2_one_hot_freq.data")
cnt = len(ch)
V = 10000

print("cnt = " + str(cnt))
print("V = " + str(V))

from src.EmbeddingGloVe import genX as genX
from src.EmbeddingGloVe import Glove as Glove

pathBase = 'dataset/chat_corpus/clean_chat_corpus/'
paths = [
    "chatterbot.tsv",
    "douban_single_turn.tsv",
    "ptt.tsv",
    "qingyun.tsv",
    "subtitle.tsv",
    "tieba.tsv",
    "weibo.tsv"
]
fullPaths = [pathBase + path for path in paths]

import numpy as np
import random

from src.GeneratorLSTMv12 import LSTMGAN

embDim = 128
hidDim = 512
conDim = 150

lstm = LSTMGAN(embDim, hidDim, conDim, V)
lstm.load("results/V12-LSTM-highflu-loss-S000.model")

def Gen(query):
    query = [id[i] if id[i] < V else 0 for i in query]
    gen = lstm.eval(query, True, 30)
    res = ''
    for c in gen:
        res += ch[c] if c > 1 else ''
    print("Gen Argmax %d > %s" % (len(res), res))

    gen, item = lstm.eval(query, True, 30, True)
    res = ''
    for c in gen:
        res += ch[c] if c > 1 else ''
    print("Gen Softmax %d > %s" % (len(res), res))
    print("Score Correlation > %s\nScore Fluency > %s" % lstm.score(query, item))
    print("------------------------------")

while True:
    cur = input("Query> ")
    if len(cur) and cur[0] == '!' :
        code = cur.split()
        if (code[0] == '!exit'):
            break
        elif (code[0] == '!code'):
            print("  Res> " + str(id[code[1]]))
    else:
        Gen(cur)
    